from sqlalchemy import *
from sqlalchemy.orm import *
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm.collections import attribute_mapped_collection
from datetime import datetime
import os

def make_pheno_table(meta, schema, name='pheno'):
    pheno_table = Table(
        name, meta,
        Column('patientid', String(60), primary_key=True),
        Column('famid', String(60), nullable=True),
        Column('sex_id',  CHAR(1), ForeignKey(schema+'.sex.val', onupdate='CASCADE', ondelete='CASCADE'), index=True),
        Column('race_id',  String(60), ForeignKey(schema+'.race.val', onupdate='CASCADE', ondelete='CASCADE'), index=True),
        Column('phenotype', SmallInteger),
        schema = schema,
        )
    return pheno_table

def make_idlink_table(meta, schema, name='idlink'):
    idlink_table = Table(
        name, meta,
        Column('id', Integer, Sequence('idlink_id_seq', schema='seq'), primary_key=True),
        Column('expid', String(60), nullable=False, unique=True, index=True),
        Column('patientid',  String(60), ForeignKey(schema+'.pheno.patientid', onupdate='CASCADE', ondelete='CASCADE'), nullable=False, index=True),
        Column('studyid', String(60), nullable=False, index=True),
        Column('sampleid', String(60), nullable=False, index=True),
        schema = schema,
        )
    return idlink_table

# Allow M (male), F (female), U (unknown).
def make_sex_table(meta, schema, name='sex'):
    sex_table = Table(
        name, meta,
        Column('val', CHAR(1), primary_key=True),
        Column('code', Integer, unique=True, nullable=False),
        Column('genabel_code', String(2), unique=True, nullable=False),
        schema = schema,
        )
    return sex_table

def make_race_table(meta, schema, name='race'):
    race_table = Table(
        name, meta,
        Column('val', String(60), primary_key=True),
        schema = schema,
        )
    return race_table

def make_anno_affy6_table(meta, schema, name='anno'):
    anno_table = Table(
        name, meta,
        Column('id', Integer, primary_key=True),
        Column('fid', String(35), nullable=False, unique=True),
        Column('rsid', String(20),  index=True, nullable=False),
        Column('chromosome', Integer, ForeignKey(schema+'.chromo.id', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('location', Integer, index=True, nullable=False),
        Column('allelea_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('alleleb_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        schema = schema,
        )
    return anno_table

def make_anno_illumina_table(meta, schema, name='anno'):
    anno_table = Table(
        'anno', meta,
        Column('id', Integer, primary_key=True),
        Column('fid', String(35), nullable=False, unique=True),
        Column('rsid', String(20),  index=True, nullable=False),
        Column('chromosome', Integer, ForeignKey(schema+'.chromo.id', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('location', Integer, index=True, nullable=False),
        Column('allelea_forward_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('alleleb_forward_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('allelea_top_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        Column('alleleb_top_id', String(1), ForeignKey(schema+'.allele.val', onupdate='CASCADE', ondelete='CASCADE'), index=True, nullable=False),
        schema = schema,
        )
    return anno_table

def make_geno_table(meta, schema, name='geno'):
    geno_table = Table(
        name, meta,
        Column('idlink_id', Integer, ForeignKey(schema+'.idlink.id', onupdate='CASCADE', ondelete='CASCADE'), nullable=False, primary_key=True),
        Column('anno_id', Integer, ForeignKey(schema+'.anno.id', onupdate='CASCADE', ondelete='CASCADE'), nullable=False, primary_key=True),
        Column('snpval_id', Integer, ForeignKey(schema+'.snpval.val', onupdate='CASCADE', ondelete='CASCADE'), nullable=False),
        schema = schema,
        )
    return geno_table

def make_geno_idlink_table(meta, schema, name="geno_idlink"):
    geno_idlink_table = Table(
        name, meta,
        Column('id',  Integer, ForeignKey(schema+'.idlink.id', onupdate='CASCADE', ondelete='CASCADE'), nullable=False, primary_key=True),
        schema = schema,
        )
    return geno_idlink_table

def make_geno_shard_table(meta, schema, id):
    name = 'geno'+str(id)
    geno_table = Table(
        name, meta,
        Column('idlink_id', Integer, CheckConstraint('idlink_id = %s'%id), ForeignKey(schema+'.geno_idlink.id', onupdate='CASCADE', ondelete='CASCADE'), nullable=False),
        Column('anno_id', Integer, ForeignKey(schema+'.anno.id', onupdate='CASCADE', ondelete='CASCADE'), nullable=False, primary_key=True),
        Column('snpval_id', Integer, ForeignKey(schema+'.snpval.val', onupdate='CASCADE', ondelete='CASCADE'), nullable=False),
        schema = schema,
        )
    return geno_table

def make_chromo_table(meta, schema, name='chromo'):
    chromo_table = Table(
        'chromo', meta,
        Column('id', Integer, primary_key=True),
        Column('name', String(2), nullable=False, index=True, unique=True),
        schema = schema,
        )
    return chromo_table

def make_allele_table(meta, schema, name='allele'):
    allele_table = Table(
        'allele', meta,
        Column('val', String(1), primary_key=True),
        schema = schema,
        )
    return allele_table

def make_snpval_table(meta, schema, name='snpval'):
    snpval_table = Table(
        'snpval', meta,
        Column('val', Integer, primary_key=True),
        schema = schema,
        )
    return snpval_table

class Pheno(object):
    def __init__(self, patientid, famid=None, sex=None, race=None, phenotype=None):
        self.patientid = patientid
        self.famid = famid
        self.sex = sex
        self.race = race
        self.phenotype = phenotype
    def __repr__(self):
        return '<Pheno %s>'%self.patientid

class Idlink(object):
    def __init__(self, expid, patientid, studyid, sampleid):
        self.expid = expid
        self.patientid = patientid
        self.studyid = studyid
        self.sampleid = sampleid
    def __repr__(self):
        return '<Idlink %s>'%self.expid

class Sex(object):
    def __init__(self, val, code=None, genabel_code=None):
        self.val = val
        self.code = code
        self.genabel_code = genabel_code
    def __repr__(self):
        return '<Sex %s>'%self.val

class Race(object):
    def __init__(self, val):
        self.val = val
    def __repr__(self):
        return '<Race %s>'%self.val

class Anno_affy6_table(object):
    def __init__(self, fid, rsid, chromo, location, alleleA, alleleB):
        self.fid = fid
        self.rsid = rsid
        self.chromo = chromo
        self.location = location
        self.alleleA = alleleA
        self.alleleB = alleleB
    def __repr__(self):
        return '<ANNO %s>'%self.fid

def Anno_affy6_mapper(Anno_table_class, anno_table, chromo_table, allele_table):
    mapper(Anno_table_class, anno_table, properties={
                                         'chromo':relation(Chromo, backref='annos_chromo', primaryjoin=anno_table.c.chromosome==chromo_table.c.id),
                                         'alleleA':relation(Allele, backref='annos_alleleA', primaryjoin=anno_table.c.allelea_id==allele_table.c.val),
                                         'alleleB':relation(Allele, backref='annos_alleleB', primaryjoin=anno_table.c.alleleb_id==allele_table.c.val),
                                         })

class Anno_illumina_table(object):
    def __init__(self, fid, rsid, chromo, location, alleleA_forward, alleleB_forward, alleleA_top, alleleB_top):
        self.fid = fid
        self.rsid = rsid
        self.chromo = chromo
        self.location = location
        self.alleleA_forward = alleleA_forward
        self.alleleB_forward = alleleB_forward
        self.alleleA_top = alleleA_top
        self.alleleB_top = alleleB_top
    def __repr__(self):
        return '<ANNO %s>'%self.fid

def Anno_illumina_mapper(Anno_table_class, anno_table, chromo_table, allele_table):
    mapper(Anno_table_class, anno_table, properties={
                                         'chromo':relation(Chromo, backref='annos_chromo', primaryjoin=anno_table.c.chromosome==chromo_table.c.id),
                                         'alleleA_forward':relation(Allele, backref='annos_alleleA', primaryjoin=anno_table.c.allelea_forward_id==allele_table.c.val),
                                         'alleleB_forward':relation(Allele, backref='annos_alleleB', primaryjoin=anno_table.c.alleleb_forward_id==allele_table.c.val),
                                         })

class Geno_table(object):
    def __init__(self, expid, anno, snpval):
        self.expid = expid
        self.anno = anno
        self.snpval = snpval
    def __repr__(self):
        return '<Geno %s>'%self.snpval

class Chromo(object):
    def __init__(self, i, name):
        self.id = i
        self.name = name
    def __repr__(self):
        return '<Chromo %s>'%self.id

class Allele(object):
    def __init__(self, val):
        self.val = val
    def __repr__(self):
        return '<Allele %s>'%self.val

class Snpval(object):
    def __init__(self, val):
        self.val = val
    def __repr__(self):
        return '<Snpval %s>'%self.val

def make_pheno_all(schema):
        from sqlalchemy import MetaData
        meta = MetaData()
        pheno_table = make_pheno_table(meta, schema)
        sex_table = make_sex_table(meta, schema)
        race_table = make_race_table(meta, schema)
        mapper(Pheno, pheno_table, properties=
               {'sex':relation(Sex, backref='patients'),
                'race':relation(Race, backref='patients'),
                })
        mapper(Sex, sex_table)
        mapper(Race, race_table)
        table_dict = {'metadata': meta, 'pheno_table':pheno_table, 'sex_table':sex_table, 'race_table':race_table, 'Pheno':Pheno, 'Sex':Sex, 'Race':Race}
        return table_dict


class Tables(object):
    def __init__(self, make_anno_table, Anno_table_class, Anno_mapper):
        self.make_anno_table = make_anno_table
        self.Anno_table_class = Anno_table_class
        self.Anno_mapper = Anno_mapper

    def make_tables(self, schema):
        from sqlalchemy import MetaData
        meta = MetaData()
        pheno_table = make_pheno_table(meta, schema)
        idlink_table = make_idlink_table(meta, schema)
        sex_table = make_sex_table(meta, schema)
        race_table = make_race_table(meta, schema)
        anno_table = self.make_anno_table(meta, schema)
        geno_table = make_geno_table(meta, schema)
        chromo_table = make_chromo_table(meta, schema)
        allele_table = make_allele_table(meta, schema)
        snpval_table = make_snpval_table(meta, schema)
        mapper(Pheno, pheno_table, properties=
               {'sex':relation(Sex, backref='patients'),
                'race':relation(Race, backref='patients'),
                })
        mapper(Idlink, idlink_table, properties={})
        mapper(Sex, sex_table)
        mapper(Race, race_table)
        self.Anno_mapper(self.Anno_table_class, anno_table, chromo_table, allele_table)
        mapper(Geno_table, geno_table, properties={'chip':relation(Idlink, backref='genos'),
                                             'anno':relation(self.Anno_table_class, backref='genos'),
                                             'snpval':relation(Snpval, uselist=False, backref='geno'),
                                             })
        mapper(Chromo, chromo_table)
        mapper(Allele, allele_table)
        mapper(Snpval, snpval_table)
        table_dict = {'metadata': meta, 'pheno_table':pheno_table, 'idlink_table':idlink_table, 'sex_table':sex_table, 'race_table':race_table, 'anno_table':anno_table, 'geno_table':geno_table, 'chromo_table':chromo_table, 'allele_table':allele_table, 'snpval_table':snpval_table, 'Pheno':Pheno, 'Idlink':Idlink, 'Sex':Sex, 'Race':Race, 'Anno':self.Anno_table_class, 'Geno_table':Geno_table, 'Chromo':Chromo, 'Allele':Allele, 'Snpval':Snpval}
        return table_dict

    def make_tables_shard(self, schema):
        """The shard table dict does not include the Geno tables"""
        from sqlalchemy import MetaData
        meta = MetaData()
        pheno_table = make_pheno_table(meta, schema)
        idlink_table = make_idlink_table(meta, schema)
        sex_table = make_sex_table(meta, schema)
        race_table = make_race_table(meta, schema)
        anno_table = self.make_anno_table(meta, schema)
        chromo_table = make_chromo_table(meta, schema)
        allele_table = make_allele_table(meta, schema)
        snpval_table = make_snpval_table(meta, schema)
        mapper(Pheno, pheno_table, properties=
               {'sex':relation(Sex, backref='patients'),
                'race':relation(Race, backref='patients'),
                })
        mapper(Idlink, idlink_table, properties={})
        mapper(Sex, sex_table)
        mapper(Race, race_table)
        self.Anno_mapper(self.Anno_table_class, anno_table, chromo_table, allele_table)
        mapper(Chromo, chromo_table)
        mapper(Allele, allele_table)
        mapper(Snpval, snpval_table)
        table_dict = {'metadata': meta, 'pheno_table':pheno_table, 'idlink_table':idlink_table, 'sex_table':sex_table, 'race_table':race_table, 'anno_table':anno_table, 'chromo_table':chromo_table, 'allele_table':allele_table, 'snpval_table':snpval_table, 'Pheno':Pheno, 'Idlink':Idlink, 'Sex':Sex, 'Race':Race, 'Anno':self.Anno_table_class, 'Chromo':Chromo, 'Allele':Allele, 'Snpval':Snpval}
        return table_dict

class Tables_affy6(Tables):
    pass

class Tables_illumina(Tables):
    def __init__(self, make_anno_illumina_table, Anno_illumina, Anno_illumina_mapper):
        Tables.__init__(self, make_anno_illumina_table, Anno_illumina, Anno_illumina_mapper)

a = Tables_affy6(make_anno_affy6_table, Anno_affy6_table, Anno_affy6_mapper)
i = Tables_illumina(make_anno_illumina_table, Anno_illumina_table, Anno_illumina_mapper)
make_tables_affy6 = a.make_tables
make_tables_shard_affy6 = a.make_tables_shard
make_tables_illumina = i.make_tables
make_tables_shard_illumina = i.make_tables_shard
